// Copyright 2022 jeff.li. and/or its affiliates.


#include <tbir/runtime/c_runtime_api.h>
#include <tbir/runtime/device_api.h>
#include <tbir/runtime/dlpack.h>
#include <tbir/runtime/generic/generic_funcs.h>
#include <tbir/runtime/registry.h>

namespace tbir::runtime {

    // set device api
    TBIR_REGISTER_GLOBAL("runtime.SetDevice").set_body([](PyArgs args) -> RTValue {
        TbirDevice device;
        device.device_type = static_cast<DLDeviceType>(args[0].As<int64_t>());
        device.device_id = args[1].As<int64_t>();
        DeviceAPI::Get(device)->SetDevice(device);
        return None;
    });

    // set device api
    TBIR_REGISTER_GLOBAL("runtime.GetDeviceAttr").set_body([](PyArgs args) -> RTValue {
        TbirDevice ctx;
        ctx.device_type = static_cast<DLDeviceType>(args[0].As<int64_t>());
        ctx.device_id = args[1].As<int64_t>();

        RTValue ret;
        DeviceAttrKind kind = static_cast<DeviceAttrKind>(args[2].As<int64_t>());
        if (kind == kExist) {
            DeviceAPI *api = DeviceAPI::Get(ctx, true);
            if (api != nullptr) {
                api->GetAttr(ctx, kind, &ret);
            } else {
                ret = 0;
            }
        } else {
            DeviceAPI::Get(ctx)->GetAttr(ctx, kind, &ret);
        }
        return ret;
    });

    TBIR_REGISTER_GLOBAL("runtime.TbirSetCurrentThreadStream")
    .set_body_typed(TbirSetCurrentThreadStream);

    // create stream
    TBIR_REGISTER_GLOBAL("runtime.DefaultStream").set_body([](PyArgs args) -> RTValue {
        MXCHECK(args.size() == 1) << "cuda_module_default_stream expect 1 args, bug get " << args.size();
        MXCHECK(args[0].type_code() == TypeIndex::kRuntimeInteger)
            << "Create Stream first arg must be integer. ";
        int device_id = args[0].As<int64_t>();
        MXCHECK(device_id >= 0) << "Device Id must be equal or greater than zeros .";

        return kernel_cuda_module_create_stream(device_id);
    });

    // create stream
    TBIR_REGISTER_GLOBAL("runtime.CreateStream").set_body([](PyArgs args) -> RTValue {
        MXCHECK(args.size() == 1) << "cuda_module_create_stream expect 1 args, bug get " << args.size();
        MXCHECK(args[0].type_code() == TypeIndex::kRuntimeInteger)
            << "Create Stream first arg must be integer. ";
        int device_id = args[0].As<int64_t>();
        MXCHECK(device_id >= 0) << "Device Id must be equal or greater than zeros .";

        return kernel_cuda_module_create_stream(device_id);
    });

    // StreamSync
    TBIR_REGISTER_GLOBAL("runtime.StreamSync").set_body([](PyArgs args) -> RTValue {
        MXCHECK(args.size() == 2) << "StreamSync expect 2 args, bug get " << args.size();
        RTView opaq = args[0].As<RTView>();
        int device_id = args[1].As<int64_t>();
        MXCHECK(device_id >= 0) << "Device Id must be equal or greater than zeros .";

        kernel_cuda_module_stream_sync(opaq, device_id);
        return None;
    });

    // StreamSync
    TBIR_REGISTER_GLOBAL("runtime.CurrentThreadStreamSync").set_body([](PyArgs args) -> RTValue {
        MXCHECK(args.size() == 2) << "CurrentThreadStreamSync expect 2 args, bug get " << args.size();
        MXCHECK(args[0].type_code() == TypeIndex::kRuntimeInteger)
            << "CurrentThreadStreamSync first arg must be integer. ";
        MXCHECK(args[1].type_code() == TypeIndex::kRuntimeInteger)
            << "CurrentThreadStreamSync first arg must be integer. ";
        int64_t device_type = args[0].AsNoCheck<int64_t>();
        int64_t device_id = args[1].AsNoCheck<int64_t>();
        TbirDevice device{DLDeviceType(device_type), DLDeviceType(device_id)};
        DeviceAPI::Get(device)->CurrentThreadStreamSync(device);
        return None;
    });

}  // namespace tbir::runtime
